{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from math import log10\n",
    "import jieba\n",
    "import Spider\n",
    "\n",
    "def highlight(item, query: str, side_len: int = 12) -> str:\n",
    "    positions = list()\n",
    "    query_words = list(jieba.cut(query))  # 把生成器强制转换为列表\n",
    "    i = 0\n",
    "    content_lower = item[2].lower()\n",
    "    word_start_map = list()\n",
    "    word_end_map = list()\n",
    "    last_word_end = -1\n",
    "    len_content_lower = len(content_lower)\n",
    "    segments = list()\n",
    "    for keyword in query_words:\n",
    "        idx = content_lower.find(keyword.lower())\n",
    "        positions.append(idx)\n",
    "    for keyword in jieba.cut(content_lower):\n",
    "        # 用于实现提取摘要时“整词切分”，避免出现截取摘要时首尾的词被截断\n",
    "        current_word_start = last_word_end + 1\n",
    "        current_word_end = current_word_start + len(keyword) - 1\n",
    "        for _ in range(current_word_start, current_word_end+1):\n",
    "            word_start_map.append(current_word_start)\n",
    "            word_end_map.append(current_word_end)\n",
    "        last_word_end = current_word_end\n",
    "    positions.sort()\n",
    "    while i < len(positions):\n",
    "        start_pos = max(positions[i] - side_len, 0)\n",
    "        end_pos = min(positions[i] + side_len, len_content_lower-1)\n",
    "        # 用于实现合并相邻且有部分重合的摘要\n",
    "        while (i < len(positions) - 1) and (positions[i+1] - positions[i] < side_len*2):\n",
    "            end_pos = min(positions[i+1] + side_len, len_content_lower-1)\n",
    "            i += 1\n",
    "        start_ellipsis = '...' if start_pos > 0 else ''\n",
    "        end_ellipsis = '...' if end_pos < len_content_lower else ''\n",
    "        segments.append(start_ellipsis + item[2][word_start_map[start_pos]: word_end_map[end_pos]] + end_ellipsis)\n",
    "        i += 1\n",
    "    result = text = item[1] + '\\n' + ''.join(segments)\n",
    "    text_lower = text.lower()\n",
    "    for keyword in query_words:\n",
    "        # 高亮部分\n",
    "        idx = text_lower.find(keyword.lower())\n",
    "        if idx >= 0:\n",
    "            ori_word = text[idx:idx+(len(keyword))]\n",
    "            result = result.replace(ori_word, f'<span style=\"color:red\";>{ori_word}</span>')\n",
    "    return result\n",
    "\n",
    "\n",
    "class MySearcherC12V0:\n",
    "    \"\"\"\n",
    "    第十一次课升级的搜索类版本：\n",
    "    改善文档频和文档长度加权的影响\n",
    "    改善IDF权值\n",
    "    采用BM25打分函数\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        self.docs = list()  # 所有文档原始数据\n",
    "        self.load_data()\n",
    "        self.cache = dict()\n",
    "        self.vocab = set()  # 索引词表\n",
    "        self.lower_preprocess()\n",
    "        jieba.load_userdict('./dict.txt')\n",
    "        self.df = dict()\n",
    "        self.avg_dl = 0\n",
    "        self.build_cache()\n",
    "\n",
    "    def load_data(self, data_file_name='./news_list.pkl'):\n",
    "        if os.path.exists(data_file_name):\n",
    "            self.docs = Spider.pickle_load(data_file_name)\n",
    "        else:\n",
    "            Spider.pickle_save(data_file_name)\n",
    "            self.docs = Spider.pickle_load(data_file_name)\n",
    "\n",
    "    def search(self, query):\n",
    "        result = None\n",
    "        for keyword in jieba.cut(query.lower()):\n",
    "            if keyword in self.cache:\n",
    "                if result is None:\n",
    "                    result = self.cache[keyword]\n",
    "                else:\n",
    "                    result = result & self.cache[keyword]\n",
    "            else:\n",
    "                result = set()\n",
    "                break\n",
    "        if result is None:\n",
    "            result = set()\n",
    "        sorted_result = self.rank(query, result)\n",
    "        return sorted_result\n",
    "\n",
    "    def rank(self, query, result_set):\n",
    "        result = list()\n",
    "        for doc_id in result_set:\n",
    "            result.append([doc_id, self.score(self.docs[doc_id], query)])\n",
    "        result.sort(key=lambda x: x[1], reverse=True)\n",
    "        return result\n",
    "\n",
    "    def render_search_result(self, query):\n",
    "        \"\"\"\n",
    "        返回带有高亮和摘要的查询结果\n",
    "        \"\"\"\n",
    "        count = 0\n",
    "        result = ''\n",
    "        for item in self.search(query):\n",
    "            count += 1\n",
    "            result += f'{count}[{item[1]}] {highlight(self.docs[item[0]], query)}\\n'\n",
    "        return result\n",
    "\n",
    "    def score(self, item, query, k1 = 2, b = 0.75):\n",
    "        score = 0\n",
    "        for keyword in jieba.cut(query.lower()):\n",
    "            f = item[2].lower().count(keyword.lower())\n",
    "            dl = len(item[2])\n",
    "            tf = (f * (k1 + 1)) / (f + k1 * (1 - b + b * (dl / self.avg_dl)))\n",
    "            idf = log10((len(self.docs) - self.df[keyword] + 0.5) / (self.df[keyword] + 0.5))\n",
    "            score += tf * idf\n",
    "        return score\n",
    "\n",
    "    def build_cache(self):\n",
    "        \"\"\"\n",
    "        用分词（用文档过滤词库）的方式构建索引\n",
    "        \"\"\"\n",
    "        doc_id = 0\n",
    "        doc_length_sum = 0\n",
    "        for doc in self.docs:\n",
    "            doc_word_set = set()\n",
    "            doc_length_sum += len(doc[3])\n",
    "            for word in jieba.cut_for_search(doc[3]):\n",
    "                if word not in doc_word_set:\n",
    "                    result_item = doc_id\n",
    "                    if word not in self.cache:\n",
    "                        self.cache[word] = {result_item}\n",
    "                    else:\n",
    "                        self.cache[word].add(result_item)\n",
    "                    self.vocab.add(word)\n",
    "                    doc_word_set.add(word)\n",
    "                    if word in self.df:\n",
    "                        self.df[word] += 1\n",
    "                    else:\n",
    "                        self.df[word] = 1\n",
    "            doc_id += 1\n",
    "        self.avg_dl = doc_length_sum / len(self.docs)\n",
    "\n",
    "    def lower_preprocess(self):\n",
    "        for doc_id in range(len(self.docs)):\n",
    "            self.docs[doc_id].append(\n",
    "                (self.docs[doc_id][1] + ' ' + self.docs[doc_id][2]).lower())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_phrase_match(phrase):\n",
    "    result = {}\n",
    "    word_id = 0\n",
    "    for word in jieba.cut(phrase):\n",
    "        if word not in self.cache:\n",
    "            result = {}\n",
    "            break\n",
    "        if word_id == 0:\n",
    "            result = self.cache.get(word, {})\n",
    "        else:\n",
    "            if len(result)  == 0:\n",
    "                break\n",
    "            result = result & self.get(word, {})\n",
    "        word_id += 1\n",
    "    return result"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def conv_query(query) -> str:\n",
    "    \"\"\"\n",
    "    将布尔查询转换成集合运算表达式\n",
    "    :param query: 查询语句\n",
    "    :return: 返回处理后的查询串\n",
    "    \"\"\"\n",
    "    query += ' '  # 在查询语句末尾加空格，确保最后一项也能正常添加\n",
    "    query_len = len(query)\n",
    "    idx = 0\n",
    "    result_parts = list()\n",
    "    cache = ''\n",
    "    while idx < query_len:\n",
    "        if query[idx] in ('(', ')', ' ') or idx + 1 == query_len:\n",
    "            if cache != '':\n",
    "                # 每当遇到括号或空格，把之前的单字片段添加到result_parts\n",
    "                if cache == 'and' or cache == 'AND':\n",
    "                    result_parts.append('&')  # 若为and则添加&运算符\n",
    "                elif cache == 'or' or cache == 'OR':\n",
    "                    result_parts.append('|')\n",
    "                elif cache == 'not' or cache == 'NOT':\n",
    "                    result_parts.append('-')\n",
    "                else:\n",
    "                    result_parts.append(f'get_phrase_match(\"{cache}\")')  # 如为查询项则直接添加\n",
    "                cache = ''\n",
    "        else:\n",
    "            cache += query[idx]\n",
    "        idx += 1\n",
    "\n",
    "    return ''.join(result_parts)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}